Add BLAS dispatches and Dot-based lowering for the JIT backend - #2406
Add BLAS dispatches and Dot-based lowering for the JIT backend#2406jessegrabowski wants to merge 10 commits into
Conversation
Introduce the singledispatch c_funcify registry returning detached CImpl implementations, resolve CLinker through it, and route OpWiseCLinker, the VM, and DebugMode to the dispatched C thunk with a Python fallback.
jessegrabowski
changed the title
The deleted make_c_gemv_destructive also duplicated a shared AllocEmpty so each Gemv could destroy its own buffer, so test_multiple_inplace fails until the generic replacement lands two commits later.
Without this the rewrite that introduces Ger regresses against the elemwise it replaces, since numba falls back to object mode for it.
Copying the accumulator in and letting BLAS scale it on top touches the output twice, which cost more than the elemwise dot-and-add these ops replace.
97bf342 to
5ac9cb9
Compare
| return batched_dot | ||
|
|
||
|
|
||
| @jax_funcify.register(Gemm) |
There was a problem hiding this comment.
this is an argument to not bother with these ops in jax, like we don't bother with Fusion/Inplace?
| return expm, cache_version | ||
|
|
||
|
|
||
| def _gemm(A, B, C, transa=False, transb=False, alpha=1.0, beta=0.0): |
There was a problem hiding this comment.
there was a big slowdown in gemv(?) C-code with negative strides, where a copy could be avoided. we may want to do the same trick for numba
|
|
||
|
|
||
| @overload(_ger) | ||
| def _ger_impl(alpha, x, y, A): |
There was a problem hiding this comment.
isn't ger the more useless one, compared to gemv?
| b = beta.item() | ||
| if b == 1.0: | ||
| out += Z | ||
| elif b != 0.0: | ||
| out += b * Z |
There was a problem hiding this comment.
isn't this defeat the point of the scalar/mul fusion of gemm?
| def dot(x, y, out=None): | ||
| if out is None: | ||
| out = np.empty((x.shape[0], y.shape[1]), dtype=numba_dot_dtype) | ||
| return _gemm(x, y, out, False, False, 1.0, 0.0) |
There was a problem hiding this comment.
I'd be surprised if numba doesn't emit blas for np.dot already
| if numba_dot_dtype in _GEMM_DTYPES: | ||
| # `gemm` reads each operand's memory order as a transpose flag, so an | ||
| # operand that reaches here transposed costs nothing, where `np.dot` would | ||
| # have to be handed a contiguous copy of it. |
| return destroy_dependencies | ||
|
|
||
|
|
||
| def get_static_scalar(node: Apply | None, input_index: int) -> float | None: |
There was a problem hiding this comment.
I don't think this helper earns its keep, the logic to make it general is more complex than inlining it
| ) | ||
|
|
||
|
|
||
| @node_rewriter([AllocEmpty]) |
There was a problem hiding this comment.
should apply to alloc of zeros as well
| destroyers = dict.fromkeys( | ||
| client | ||
| for client, input_index in clients[1:] | ||
| if not isinstance(client.op, Output) |
There was a problem hiding this comment.
the refactor and bootstrap code for the c impl dispatcher looks neat.
OTOH I'm concerned IFF we are making the blas pipeline run by default in jax and numba? For jax it reads just like rewrite overhead, since we end up emitting the naive code. For numba I'd need a more exhaustive reproducible benchmark than "it speeds up pytensor-ml by 10%", as this is a fundamental change that touches most graphs we work with. Or proving that np.dot always lowers to blas by numba anyway and we are just skipping some indirection (which ones?).
2 participants
All of this measures netural against pymc-model-catalog. I saw 10-15% speedup in the backward pass of linear layers in pytensor-ml. Rewriting graphs of the form
a + B @ Cinto GEMM isn't perfect, it still requires the user to put parenthesis. We can try to tune it up in follow-up work, but it's not super clear to me it's an obvious win.